#include <cstring>
#include <fstream>
#include <iostream>
#include <algorithm>
#include "ti_3D.h"
#include "tplate.h"
#include "taxonomy_tree.h"

using namespace std;

void init_sim_grid_3D(target_grid_3D * grid, int length, int width,
                      int height, int max_scale, int NR_cate, int N_sample, int *target)
{
    grid->max_scale = max_scale;
    grid->length = new int[max_scale];
    grid->width = new int[max_scale];
    grid->height = new int[max_scale];
    grid->length[0] = length;
    grid->width[0] = width;
    grid->height[0] = height;
    grid->NR_cate = NR_cate;
    grid->hard_data = new int[length * width * height];
    memcpy(grid->hard_data, target, sizeof(int) * length * width * height);
    grid->sample = new target_grid_DATA[max_scale];
    for (int g = 0; g < max_scale; g++)
    {
        int cur_length = grid->length[g] = (length >> g);
        int cur_width = grid->width[g] = (width >> g);
        int cur_height = grid->height[g] = (height >> g);
        // unit size
        int unit_size = (1 << g);
        target_grid_DATA *cur_sample = grid->sample + g;
        cur_sample->category =
            new int[cur_length * cur_width * cur_height * NR_cate];
        cur_sample->prob =
            new float[cur_length * cur_width * cur_height * NR_cate];
        memset(cur_sample->category, 0,
               sizeof(int) * cur_length * cur_width * cur_height *
               NR_cate);
        memset(cur_sample->prob, 0,
               sizeof(int) * cur_length * cur_width * cur_height *
               NR_cate);
        // ensample to higher scale grid
        for (int x = 0; x < length; x++)
        {
            for (int y = 0; y < width; y++)
            {
                for (int z = 0; z < height; z++)
                {
                    int cur_cate =
                        grid->hard_data[x * width * height + y * height +
                                        z];
                    if (-1 == cur_cate)
                        continue;
                    cur_sample->category[(x / unit_size) *
                                         cur_width * cur_height * NR_cate +
                                         (y / unit_size) * cur_height *
                                         NR_cate +
                                         (z / unit_size) * NR_cate +
                                         cur_cate] += 1;
                }
            }
        }
    }
    return;
}

int *zoom_in_3D(int *orig_cimg, int orig_xsize, int orig_ysize,
                int orig_zsize, int scale, int NR_cate)
{
    int G = 1 << scale;
    int xsize = orig_xsize / G;
    int ysize = orig_ysize / G;
    int zsize = orig_zsize / G;
    int *zoomed = new int[xsize * ysize * zsize * NR_cate];
    memset(zoomed, 0, xsize * ysize * zsize * NR_cate * sizeof(int));
    for (int origi = 0; origi < orig_xsize; origi++)
    {
        int i = origi / G;
        if (i >= xsize)
            continue;
        for (int origj = 0; origj < orig_ysize; origj++)
        {
            int j = origj / G;
            if (j >= ysize)
                continue;
            for (int origk = 0; origk < orig_zsize; origk++)
            {
                int k = origk / G;
                if (k >= zsize)
                    continue;
                zoomed[i * ysize * zsize * NR_cate + j * zsize * NR_cate +
                       k * NR_cate +
                       orig_cimg[origi * orig_ysize * orig_zsize +
                                 origj * orig_zsize + origk]]++;
            }
        }
    }
    return zoomed;
}

void build_zoomed_images_3D(int maxscale, train_grid_3D * ti)
{
    ti->max_scale = maxscale;
    for (int g = 1; g < maxscale; g++)
    {
        int G = 1 << g;
        ti->map_cate[g].length = ti->length / G;
        ti->map_cate[g].width = ti->width / G;
        ti->map_cate[g].height = ti->height / G;
        ti->map_cate[g].NR_cate = ti->NR_cate;
        ti->map_cate[g].data =
            zoom_in_3D(ti->data, ti->length, ti->width, ti->height, g,
                       ti->NR_cate);
    }
}

void read_image_3D(train_grid_3D * ti, char *file)
{
    string type;
    int tmp, NR_cate;
    NR_cate = 0;
    fstream fs;
    fs.open(file);
    fs >> ti->length;
    fs >> ti->width;
    fs >> ti->height;
    fs >> tmp;
    fs >> type;
    int size = (ti->length) * (ti->width) * ti->height;
    ti->data = new int[size];
    for (int i = 0; i < size; i++)
    {
        fs >> (ti->data)[i];
        NR_cate = max(NR_cate, ti->data[i] + 1);
    }
    fs.close();
}
